import torch
import torchvision as vision
from torchvision import transforms
from PIL import Image
import torchvision.models as models
import os
from reshape import reshape_model


class Prediction:

    def __init__(self, model=None, classes=None):
        self.model = model
        self.classes = classes
        self.is_cpu = False if torch.cuda.is_available() else True
        self.transform = transforms.Compose([
            transforms.Resize(256),
            transforms.CenterCrop(224),
            transforms.ToTensor(),
            transforms.Normalize(
                mean=[0.485, 0.456, 0.406],
                std=[0.229, 0.224, 0.225]
            )
        ])
        os.environ["TORCH_HOME"] = "models/pretrained_models"

    def predict(self, image):
        if self.model is None or self.classes is None:
            return [("请先选择模型", 0.0)]
        if isinstance(image, str):
            img = Image.open(image).convert('RGB')
        else:
            img = image
        img_t = self.transform(img)
        batch_t = torch.unsqueeze(img_t, 0)
        self.model.eval()
        out = self.model(batch_t)
        # Forth, print the top 5 classes predicted by the model
        _, indices = torch.sort(out, descending=True)
        percentage = torch.nn.functional.softmax(out, dim=1)[0] * 100
        pre_number = 1 if len(self.classes) < 15 else 3
        return [(self.classes[idx], percentage[idx].item()) for idx in indices[0][:pre_number]]

    def load_pretrained_model(self, model_name, classes_path):
        self.model = eval("models." + model_name + "(pretrained=True)")
        self.get_classes(classes_path)

    def get_classes(self, classes_path):
        with open(classes_path) as f:
            self.classes = [line.strip() for line in f.readlines()]

    def load_retrained_model(self, model_path, classes):
        print('loading checkpoint:  ' + model_path)
        if self.is_cpu:
            checkpoint = torch.load(model_path, map_location=torch.device('cpu'))
        else:
            checkpoint = torch.load(model_path)
        # checkpoint = torch.load(model_path)
        arch = checkpoint['arch']
        # create the model architecture
        print('using model:  ' + arch)
        model = models.__dict__[arch](pretrained=True)
        # reshape the model's output
        model = reshape_model(model, arch, checkpoint['num_classes'])

        # load the model weights
        model.load_state_dict(checkpoint['state_dict'])

        # add softmax layer
        print('adding nn.Softmax layer to model...')
        model = torch.nn.Sequential(model, torch.nn.Softmax(1))
        self.model = model
        self.get_classes(classes)


if __name__ == '__main__':
    os.environ['TORCH_HOME'] = 'models/pretrained_models'
    resnet = vision.models.resnet18(pretrained=True)
    prediction = Prediction(resnet, 'imagenet_classes.txt')
    print(prediction.predict('e:/test.jpg'))
